#include "kf.h"
#include <iostream>
namespace KalmanFilter
{
KF::KF(std::shared_ptr<DiscreteLinearSys> dis_ptr, Eigen::MatrixXd sys_covariance, Eigen::MatrixXd measure_covariance)
{
	if (dis_ptr->distur_dim() != 0)
	{
		std::cout << "\033[31m Error:KF::constructor failed ,distur_dim dim is 0 \033[0m" << std::endl;
	}
	dis_ptr_ = dis_ptr;
	state_dim_ = dis_ptr_->state_dim();
	control_dim_ = dis_ptr_->control_dim();
	out_dim_ = dis_ptr_->out_dim();

	sys_covariance_ = Eigen::MatrixXd::Identity(state_dim_, state_dim_);
	measure_covariance_ = Eigen::MatrixXd::Identity(out_dim_, out_dim_);

	state_predict_ = Eigen::VectorXd::Zero(state_dim_);
	state_cov_predict_ = Eigen::MatrixXd::Zero(state_dim_, state_dim_);
	kalman_gain_ = Eigen::MatrixXd::Zero(state_dim_, out_dim_);
	y_error_ = Eigen::VectorXd::Zero(out_dim_);
	state_esti_ = Eigen::VectorXd::Zero(state_dim_);
	state_cov_esti_ = Eigen::MatrixXd::Identity(state_dim_, state_dim_);

	set_sys_covariance(sys_covariance);
	set_measure_covariance(measure_covariance);
}

bool KF::set_sys_covariance(Eigen::MatrixXd sys_covariance)
{
	if (sys_covariance.rows() != state_dim_ || sys_covariance.cols() != state_dim_)
	{
		std::cout << "\033[31m Error:KF::set_sys_covariance failed ,sys_covariance dim is not right \033[0m" << std::endl;
		return false;
	}
	sys_covariance_ = sys_covariance;
	return true;
}
bool KF::set_measure_covariance(Eigen::MatrixXd measure_covariance)
{
	if (measure_covariance.rows() != out_dim_ || measure_covariance.cols() != out_dim_)
	{
		std::cout << "\033[31m Error:KF::set_measure_covariance failed ,measure_covariance dim is not right \033[0m" << std::endl;
		return false;
	}
	measure_covariance_ = measure_covariance;
	return true;
}
bool KF::set_state_esti(Eigen::VectorXd state_esti)
{
	if (state_esti.rows() != state_dim_)
	{
		std::cout << "\033[31m Error:KF::set_state_esti failed ,state_esti dim is not right \033[0m" << std::endl;
		return false;
	}
	state_esti_ = state_esti;
	return true;
}
bool KF::reset(Eigen::VectorXd state_esti, Eigen::MatrixXd state_cov_esti)
{
	bool temp_flag = true;
	temp_flag &= set_state_esti(state_esti);
	temp_flag &= set_state_cov_esti(state_cov_esti);
	if (temp_flag == false)
	{
		return false;
	}
	return true;
}
bool KF::set_state_cov_esti(Eigen::MatrixXd state_cov_esti)
{
	if (state_cov_esti.rows() != state_dim_ || state_cov_esti.cols() != state_dim_)
	{
		std::cout << "\033[31m Error:KF::set_state_cov_esti failed ,state_cov_esti dim is not right \033[0m" << std::endl;
		return false;
	}
	state_cov_esti_ = state_cov_esti;
	return true;
}
bool KF::predict(Eigen::VectorXd &control_last)
{
	if (control_last.rows() != control_dim_)
	{
		std::cout << "\033[31m Error:KF::predict failed ,control_last dim is not right \033[0m" << std::endl;
		return false;
	}
	state_predict_ = dis_ptr_->state_matrix() * state_esti_ + dis_ptr_->control_matrix() * control_last;
	state_cov_predict_ = dis_ptr_->state_matrix() * state_cov_esti_ * dis_ptr_->state_matrix().transpose() +
						 sys_covariance_;
	return true;
}
bool KF::update(Eigen::VectorXd &y)
{
	if (y.rows() != out_dim_)
	{
		std::cout << "\033[31m Error:KF::update failed ,y dim is not right \033[0m" << std::endl;
		return false;
	}
	Eigen::MatrixXd temp_matrix = dis_ptr_->out_state_matrix() * state_cov_predict_ *
									  dis_ptr_->out_state_matrix().transpose() +
								  measure_covariance_;
	kalman_gain_ = state_cov_predict_ * dis_ptr_->out_state_matrix().transpose() * temp_matrix.inverse();
	y_error_ = y - dis_ptr_->out_state_matrix() * state_predict_;
	state_esti_ = state_predict_ + kalman_gain_ * y_error_;
	state_cov_esti_ = (Eigen::MatrixXd::Identity(state_dim_, state_dim_) -
					   kalman_gain_ * dis_ptr_->out_state_matrix()) *
					  state_cov_predict_;
	return true;
}
bool KF::step(Eigen::VectorXd &control_last, Eigen::VectorXd &y)
{
	if (predict(control_last) == false)
	{
		return false;
	}
	if (update(y) == false)
	{
		return false;
	}
	return true;
}

Eigen::MatrixXd KF::sys_covariance()
{
	return sys_covariance_;
}
Eigen::MatrixXd KF::measure_covariance()
{
	return measure_covariance_;
}
Eigen::VectorXd KF::state_predict()
{
	return state_predict_;
}
Eigen::MatrixXd KF::state_cov_predict()
{
	return state_cov_predict_;
}
Eigen::MatrixXd KF::kalman_gain()
{
	return kalman_gain_;
}

Eigen::VectorXd KF::state_esti()
{
	return state_esti_;
}
Eigen::MatrixXd KF::state_cov_esti()
{
	return state_cov_esti_;
}
Eigen::VectorXd KF::y_error()
{
	return y_error_;
}
} // namespace KalmanFilter